import torch

def compute_distance_matrix(A, B):
    diagA = torch.mm(A, A.transpose(1, 0)).diag()
    diagA = diagA.unsqueeze(1).repeat(1, B.size(0))

    diagB = torch.mm(B, B.transpose(1, 0)).diag()
    diagB = diagB.unsqueeze(0).repeat(A.size(0), 1)

    AB = torch.mm(A, B.transpose(1, 0))

    D = diagA - 2 * AB + diagB

    return D


def compute_self_distance_matrix(A):
    prod = torch.mm(A, A.transpose(1, 0))
    diagA = torch.diag(prod).unsqueeze(1).repeat(1, A.size(0))

    diagB = diagA.transpose(1, 0)

    D = diagA - 2 * prod + diagB
    #print diagA
    return D


def construct_k_neighbor(D, k, bandwidth):
    Np, idx = D.topk(k=D.size(0)-k, dim=1, largest=True)
    D.scatter_(1, idx[:,:],  0.0)
    W = torch.exp(-D/ bandwidth**2)
    W[W==1.0] = 0
    W = torch.max(W, W.transpose(1, 0))
    return W


def manifold_regularization(A):

    D = compute_self_distance_matrix(A)
    W = construct_k_neighbor(D, k=5, bandwidth=2)
    W = W.cuda()
    reg = D * W
    #print D
    reg = reg.sum()
    return reg / 2